import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1)
plt.figure(figsize=[15, 9])
spr = 3
spc = 5
spn = 0
n_cluster = 4
cmap = plt.cm.get_cmap('rainbow', n_cluster)

data = np.loadtxt(r'../data/test.txt', delimiter='\t')
m = len(data)
np.random.shuffle(data)
centers = data[0:n_cluster]


def xvisualization(xtitle):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(xtitle)
    plt.scatter(data[:, 0], data[:, 1], s=4, c=min_dis_idx, cmap=cmap)
    plt.scatter(centers[:, 0], centers[:, 1], marker='x', s=32, c=range(n_cluster), cmap=cmap)


iter0 = 0
while iter0 < 7:
    iter0 += 1

    # first cycle
    dis_matrix = np.zeros([m, n_cluster])
    for i, c in enumerate(centers):
        dis_vector = np.sqrt(np.sum((data - c) ** 2, axis=1))
        dis_matrix[:, i] = dis_vector
    min_dis_idx = np.argmin(dis_matrix, axis=1)
    xvisualization(f'{iter0}th - det clusters')

    # 2nd cycle
    for i, c in enumerate(centers):
        data_in_this_culster = data[min_dis_idx == i]
        new_center_i = data_in_this_culster.mean(axis=0)
        centers[i] = new_center_i
    xvisualization(f'{iter0}th - det new centers')

plt.show()
